package available

import (
	"context"
	"google.golang.org/grpc"
	"time"
)

type TokenBucketBuilder struct {
	token chan struct{}
	close chan struct{}
}

func NewTokenBucketBuilder(capacity int, duration time.Duration) *TokenBucketBuilder {
	t := &TokenBucketBuilder{
		token: make(chan struct{}, capacity),
		close: make(chan struct{}),
	}

	ticker := time.NewTicker(duration)
	go func() {
		defer ticker.Stop()
		select {
		case <-t.close:
			return
		case <-ticker.C:
			select {
			case t.token <- struct{}{}:
			default:
			}
		}
	}()

	return t
}

func (t *TokenBucketBuilder) Build() grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
		select {
		case <-t.close:
			return handler(ctx, req)
		case <-t.token:
			return handler(ctx, req)
		case <-ctx.Done():
			return nil, ctx.Err()
		}
	}
}

func (t *TokenBucketBuilder) Close() {
	close(t.close)
}
